#!/usr/bin/env python

import cPickle
import optparse
import os

from charmicat import BadCommandLineError
import histogram

from fitter import GaussPlusPolynomialFit
import cmix

allowed_dists = "m_rs m_ws m_pipi m_kk".split()
dist_latex = dict(zip(allowed_dists, [r"$M(\left[K\pi\right]_{RS})$",
                                      r"$M(\left[K\pi\right]_{WS})$",
                                      r"$M(\pi^+\pi^-)$",
                                      r"$M(K^+K^-)$",
                                      ]))

def main():
    '''
    Usage: %prog [opts] distname input.hpt
 
    Fit the D0 peak in the `distname' distribution, where `distname' 
    can be any of `m_rs', `m_ws', `m_pipi', or `m_kk'. 
    '''

    op = optparse.OptionParser(usage=main.__doc__)
    
    op.add_option("-o", "--output", dest='output', default=None,
                  help='Write output with prefix PREFIX', metavar='PREFIX')
    op.add_option("-p", "--poly-order", dest='polyorder', type='int', default=4,
                  help='Use polynomial of order N for background', metavar='N')
    op.add_option("-t", "--title", dest='title', default=None,
                  help='Set fit title to TITLE', metavar='TITLE')
    op.add_option("--fr", dest='fit_range', nargs=2, type='float', 
                  default=(1.75, 2.0), 
                  help='Specify overall fit range as LOW HIGH', 
                  metavar='LOW HIGH')
    op.add_option("--sr", dest='signal_range', nargs=2, type='float', 
                  default=(1.84, 1.89), 
                  help='Specify signal range as LOW HIGH', metavar='LOW HIGH')
    

    opts, args = op.parse_args()

    if len(args) < 2:
        msg = "Please supply the name of the distribution to fit "
        msg += "and a histogram file as input"
        raise BadCommandLineError(msg, -1)

    distname, hpath = args[:2]

    if distname not in allowed_dists:
        msg = "Unknown distribution `{0}'; choices are \n"
        msg += str(allowed_dists)
        raise BadCommandLineError(msg , -2)
    
    hc = histogram.load(hpath)

    if distname not in hc:
        msg = "Path `{0}' does not contain histogram `{1}'"
        raise BadCommandLineError(msg.format(hpath, distname) , -2)
        
    h = hc[distname]

    if opts.output is None:
        opts.output = '{0}_{1}'.format(
            os.path.basename(hpath).rsplit('.', 1)[0], distname)

    f = GaussPlusPolynomialFit(h, opts.polyorder)
    f.fit(opts.fit_range[0], opts.fit_range[1], 
          signal_range=(opts.signal_range[0], 
                        opts.signal_range[1]))
    
    f.fit(opts.fit_range[0], opts.fit_range[1], 
          A=f.hist.max_bin_content(), 
          mu=cmix.params['D0_MASS'], 
          sigma=cmix.params['D0_SIGMA_MC'])

    if opts.title is None:
        opts.title = r"$D^0$ fit in {0}".format(dist_latex[distname])

    f.plot(name="{0}_{1}".format(opts.output, 'd0fit'), 
           title=opts.title)
    
    f.fix_signal()
    f.release_background()
    f.fit(*opts.fit_range)
    f.plot(name="{0}_{1}".format(opts.output, 'd0fit_bgrel'), 
           title=opts.title)
    f.plot(name="{0}_{1}".format(opts.output, 'd0fit_bgsub'), 
           title=opts.title, subtract_bg=True)

    with open("{0}_{1}.pkl".format(opts.output, 'd0fit_pars'), 'w') as fppars:
        cPickle.dump(f.m.values, fppars)

    return 0


if __name__ == '__main__':
    try:
        import sys
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(
                os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
            
